import pygame
import math
import sys

# 取符号函数
def sign(a):
   return (a > 0) - (a < 0)

# 行星绕日函数
def planet(m, pos_x, pos_y, vel_x, vel_y):
    # 日地坐标差
    delta_x = (sun[0] - pos_x) * k
    delta_y = (sun[1] - pos_y) * k
    # 日地距离平方
    r2 = delta_x ** 2 + delta_y ** 2

    # 日地间引力，万有引力定律
    F = G * m * M / r2
    # 日地夹角
    theta = math.acos(delta_x / r2 ** 0.5)
    # x、y 轴引力分量
    fx = abs(F * math.cos(theta)) * sign(delta_x)
    fy = abs(F * math.sin(theta)) * sign(delta_y)
    # x、y 轴加速度，牛顿第二定律
    ax = fx / m
    ay = fy / m
    # 速度变化，vt = v0 + at
    vel_x += ax * t
    vel_y += ay * t
    # 位移变化，st = s0 + vt
    pos_x += vel_x * t / k
    pos_y += vel_y * t / k
    return pos_x, pos_y, vel_x, vel_y, r2**0.5

# 定义颜色
black = 0, 0, 0
white = 255, 255, 255

m_shui = 3.3e23             # 水星质量
m_jin = 4.869e24            # 金星质量
m_e = 5.9742e24               # 地球质量
m_huo = 6.4219e23               # 火星质量
m_j = 1.9e27               # 木星质量
m_tu = 5.68e26              # 土星质量
m_tian = 8.683e25           # 天王星质量
m_hai = 1.0247e26           # 海王星质量
M = 1.9891e30               # 太阳质量

k = 1e10                     # 距离缩放参数
G = 6.67259e-17             # 万有引力常量
t = 1e8                     # 时间缩放参数

distance_e = 1.495978707e11     # 日地距离
distance_j = 7.7833e11     # 日木距离
distance_huo = 2.2794e11     # 日火距离
distance_shui = 5.791e10     # 日水距离
distance_jin = 1.082e11     # 日金距离
distance_tu = 1.4294e12     # 日土距离
distance_tian = 2.87099e12     # 日天距离
distance_hai = 4.504e12     # 日海距离

sun = (500, 300)            # 太阳坐标
# v_s = 0                   # 太阳速度

vel_xshui = 0                  # 水星初速度
vel_yshui = (G*M/distance_shui)**0.5
pos_xshui = 500 - distance_shui/k    # 水星坐标
pos_yshui = 300
shuixing = (pos_xshui, pos_yshui)

vel_xjin = 0                  # 金星初速度
vel_yjin = (G*M/distance_jin)**0.5
pos_xjin = 500 - distance_jin/k    # 坐标
pos_yjin = 300
jinxing = (pos_xjin, pos_yjin)

vel_xe = 0                  # 地球初速度
vel_ye = (G*M/distance_e)**0.5
pos_xe = 500 - distance_e/k    # 地球坐标
pos_ye = 300
earth = (pos_xe, pos_ye)

vel_xhuo = 0                  # 火星初速度
vel_yhuo = (G*M/distance_huo)**0.5
pos_xhuo = 500 - distance_huo/k    # 坐标
pos_yhuo = 300
Mars = (pos_xhuo, pos_yhuo)

vel_xj = 0                  # 木星初速度
vel_yj = (G*M/distance_j)**0.5
pos_xj = 500 - distance_j/k    # 木星坐标
pos_yj = 300
jupiter = (pos_xj, pos_yj)

vel_xtu = 0                  # 土星初速度
vel_ytu = (G*M/distance_tu)**0.5
pos_xtu = 500 - distance_tu/k    # 坐标
pos_ytu = 300
tuxing = (pos_xtu, pos_ytu)

vel_xtian = 0                  # 天王星初速度
vel_ytian = (G*M/distance_tian)**0.5
pos_xtian = 500 - distance_tian/k    # 坐标
pos_ytian = 300
tianwang = (pos_xtian, pos_ytian)

vel_xhai = 0                  # 海王星初速度
vel_yhai = (G*M/distance_hai)**0.5
pos_xhai = 500 - distance_hai/k    # 坐标
pos_yhai = 300
haiwang = (pos_xhai, pos_yhai)

# 大往左上角偏，小往右下角偏
li_xs = [(44,25), (33,25),(25,25),(32,25),(100,94),(80,60),(100,75),(100,100)]
xs_s = 2                  # 太阳像素放缩比例
xs_shui = 2                  # 水星像素放缩比例
xs_jin = 2                  # 金星像素放缩比例
xs_e = 2                   # 地球像素放缩比例
xs_huo = 2                   # 火星像素放缩比例
xs_j = 2                  # 木星像素放缩比例
xs_tu = 2                  # 土星像素放缩比例
xs_tian = 2                  # 天王星像素放缩比例
xs_hai = 2                  # 海王星像素放缩比例

pygame.init()  # 初始化
screen= pygame.display.set_mode((1000, 600))  # 创建窗口
font = pygame.font.Font(r'C:\Windows\Fonts\simkai.ttf', 30)  # 显示中文需要字体，否则可略过
text = font.render("引力弹弓", 1, white)
guiji = font.render(".", 1, white)
pygame.display.set_caption("引力弹弓模拟")
ls_e, ls_j = [], []    # 存放轨迹点
ls_pos = []

s = pygame.image.load("太阳1.png").convert_alpha()         # 太阳图片
shui = pygame.image.load("水星1.png").convert_alpha()      # 水星图片
jin = pygame.image.load("金星1.png").convert_alpha()       # 金星图片
e = pygame.image.load("earth.png").convert_alpha()       # 地球图片
Ma = pygame.image.load("火星1.png").convert_alpha()        # 火星图片
jup = pygame.image.load("jupiter.png").convert_alpha()   # 木星图片
tu = pygame.image.load("土星1.png").convert_alpha()        # 土星图片
tian = pygame.image.load("天王星1.png").convert_alpha()    # 天王星图片
hai = pygame.image.load("海王星1.png").convert_alpha()     # 海王星图片

while True:  # 主循环
    for event in pygame.event.get():
        if event.type in (pygame.QUIT,):
            sys.exit()  # 按键响应，按键后退出
    
    screen.fill(black)  # 刷新背景
    pos_sun = (sun[0] - 40 / xs_s, sun[1] - 40 / xs_s)
    screen.blit(s, pos_sun)  # 画太阳
    # 地球轨迹
    # for i_e, i_j in zip(ls_e[-300:], ls_j[-300:]):
    #     screen.blit(guiji, i_e)
    #     screen.blit(guiji, i_j)

    # for m, pos_x, pos_y, vel_x, vel_y, distance, xy, xy_change, img in zip(li_m, li_posx, li_posy, li_vx, li_vy, li_distance, li_xy, li_xs,li_img):
    #     pos_x, pos_y, vel_x, vel_y, distance = planet(m, pos_x, pos_y, vel_x, vel_y)
    #     xy = (pos_x, pos_y)
    #     xy_mid = (pos_x - xy_change[0] / 2, pos_y - xy_change[1] / 2)
    #     screen.blit(img, xy_mid)  # 画水星
    pos_xshui, pos_yshui, vel_xshui, vel_yshui, distance_shui = planet(m_shui, pos_xshui, pos_yshui, vel_xshui, vel_yshui)
    shuixing = (pos_xshui, pos_yshui)
    pos_shui = (pos_xshui - 44 / xs_shui, pos_yshui - 25 / xs_shui)
    screen.blit(shui, pos_shui)  # 画水星

    pos_xjin, pos_yjin, vel_xjin, vel_yjin, distance_jin = planet(m_jin, pos_xjin, pos_yjin, vel_xjin, vel_yjin)
    jinxing = (pos_xjin, pos_yjin)
    pos_jin = (pos_xjin - 33 / xs_jin, pos_yjin - 25 / xs_jin)
    screen.blit(jin, pos_jin)  # 画金星

    pos_xe, pos_ye, vel_xe, vel_ye, distance_e = planet(m_e, pos_xe, pos_ye, vel_xe, vel_ye)
    earth = (pos_xe, pos_ye)
    # ls_e.append(earth)
    pos_earth = (pos_xe-25/xs_e, pos_ye-25/xs_e)
    screen.blit(e, pos_earth)  # 画地球

    pos_xhuo, pos_yhuo, vel_xhuo, vel_yhuo, distance_huo = planet(m_huo, pos_xhuo, pos_yhuo, vel_xhuo, vel_yhuo)
    Mars = (pos_xhuo, pos_yhuo)
    pos_huo = (pos_xhuo - 32 / xs_huo, pos_yhuo - 25 / xs_huo)
    screen.blit(Ma, pos_huo)  # 画火星

    pos_xj, pos_yj, vel_xj, vel_yj, distance_j = planet(m_j, pos_xj, pos_yj, vel_xj, vel_yj)
    jupiter = (pos_xj, pos_yj)
    # ls_j.append(jupiter)
    pos_jupiter = (pos_xj - 100 / xs_j, pos_yj - 94 / xs_j)
    screen.blit(jup, pos_jupiter)  # 画木星

    pos_xtu, pos_ytu, vel_xtu, vel_ytu, distance_tu = planet(m_tu, pos_xtu, pos_ytu, vel_xtu, vel_ytu)
    tuxing = (pos_xtu, pos_ytu)
    pos_tu = (pos_xtu - 80 / xs_tu, pos_ytu - 60 / xs_tu)
    screen.blit(tu, pos_tu)  # 画土星

    pos_xtian, pos_ytian, vel_xtian, vel_ytian, distance_tian = planet(m_tian, pos_xtian, pos_ytian, vel_xtian, vel_ytian)
    tianwang = (pos_xtian, pos_ytian)
    pos_tian = (pos_xtian - 100 / xs_tian, pos_ytian - 75 / xs_tian)
    screen.blit(tian, pos_tian)  # 画天王星

    pos_xhai, pos_yhai, vel_xhai, vel_yhai, distance_hai = planet(m_hai, pos_xhai, pos_yhai, vel_xhai, vel_yhai)
    haiwang = (pos_xhai, pos_yhai)
    pos_hai = (pos_xhai - 100 / xs_hai, pos_yhai - 100 / xs_hai)
    screen.blit(hai, pos_hai)  # 画海王星

    v_e = '地球速度= %.1f km/s' % ((vel_xe ** 2 + vel_ye ** 2) ** 0.5)  # 速度大小
    speed_e = font.render(v_e, 1, white)
    screen.blit(speed_e, (100, 50))
    v_j = '木星速度= %.1f km/s' % ((vel_xj ** 2 + vel_yj ** 2) ** 0.5)  # 速度大小
    speed_j = font.render(v_j, 1, white)
    screen.blit(speed_j, (500, 50))
    screen.blit(font.render('日地距离%.2f×10^8 km' % (distance_e / 1e11), 1, white), (100, 20))
    screen.blit(font.render('日木距离%.2f×10^8 km' % (distance_j / 1e11), 1, white), (500, 20))
    pygame.display.update()  # 刷新
